#include "sm3.h"
#include <stdio.h>
#include <string.h>
#include <assert.h>

// ----------------------------------------------------------------------------
// SM3 Cryptographic Hash Algorithm Impl
// For bit string with len < 2^64
// ----------------------------------------------------------------------------
#define SM3_BIT_GROUP_SIZE 512
typedef uint8_t sm3_bits256[32];
typedef uint8_t sm3_bits512[64];

// SM3 Debug Output
void sm3_show_bytes(const uint8_t *bytes, uint32_t byte_len) {
    for (uint64_t i = 0; i < byte_len; ++i) {
        if (i > 0 && i % 32 == 0) {
            printf("\n");
        }
        printf("%02X", bytes[i]);
        if (i % 4 == 3) {
            printf(" ");
        }
    }
    printf("\n");
}

// SM3 Constants - default IV
const static uint8_t SM3_IV[32] = {
    0x73, 0x80, 0x16, 0x6F, 0x49, 0x14, 0xB2, 0xB9,
    0x17, 0x24, 0x42, 0xD7, 0xDA, 0x8A, 0x06, 0x00,
    0xA9, 0x6F, 0x30, 0xBC, 0x16, 0x31, 0x38, 0xAA,
    0xE3, 0x8D, 0xEE, 0x4D, 0xB0, 0xFB, 0x0E, 0x4E
};
const static uint32_t SM3_T0 = 0x79CC4519;
const static uint32_t SM3_T1 = 0x7A879D8A;

// SM3 word rotate left
uint32_t sm3_rotate_left(uint32_t x, uint32_t k) {
    k = k % 32;
    return k == 0 ? x : ((x << k) | (x >> (32 - k)));
}

// SM3 boolean functions
uint32_t sm3_ff(
    const uint32_t x,
    const uint32_t y,
    const uint32_t z,
    const int j
) {
    assert(j >= 0 && j <= 63);
    if (j <= 15) {
        return x ^ y ^ z;
    }
    return (x & y) | (x & z) | (y & z);
}

uint32_t sm3_gg(
    const uint32_t x,
    const uint32_t y,
    const uint32_t z,
    const int j
) {
    assert(j >= 0 && j <= 63);
    if (j <= 15) {
        return x ^ y ^ z;
    }
    return (x & y) | (~x & z);
}

// SM3 permutation functions
uint32_t sm3_p0(const uint32_t x) {
    return x ^ sm3_rotate_left(x, 9) ^ sm3_rotate_left(x, 17);
}

uint32_t sm3_p1(const uint32_t x) {
    return x ^ sm3_rotate_left(x, 15) ^ sm3_rotate_left(x, 23);
}

// SM3 get and set data
uint32_t sm3_data_get_uint32(const sm3_bits256 x, const int i) {
    const uint32_t b0 = x[i * 4];
    const uint32_t b1 = x[i * 4 + 1];
    const uint32_t b2 = x[i * 4 + 2];
    const uint32_t b3 = x[i * 4 + 3];
    return (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
}

void sm3_data_xor_uint32(sm3_bits256 x, const uint32_t val, const int i) {
    x[i * 4] ^= (val >> 24) & 0xFF;
    x[i * 4 + 1] ^= (val >> 16) & 0xFF;
    x[i * 4 + 2] ^= (val >> 8) & 0xFF;
    x[i * 4 + 3] ^= val & 0xFF;
}

// SM3 Compress Function
void sm3_cf(sm3_bits256 v, const sm3_bits512 msg) {
    // printf("Msg Segment:\n"); sm3_show_bytes(msg, SM3_BIT_GROUP_SIZE / 8);
    uint32_t ss1 = 0;
    uint32_t ss2 = 0;
    uint32_t tt1 = 0;
    uint32_t tt2 = 0;
    uint32_t w1[68];
    uint32_t w2[64];
    for (int j = 0; j < 16; ++j) {
        const uint32_t b0 = msg[j * 4];
        const uint32_t b1 = msg[j * 4 + 1];
        const uint32_t b2 = msg[j * 4 + 2];
        const uint32_t b3 = msg[j * 4 + 3];
        w1[j] = (b0 << 24) | (b1 << 16) | (b2 << 8) | b3;
    }
    for (int j = 16; j < 68; ++j) {
        w1[j] = sm3_p1(w1[j - 16] ^ w1[j - 9] ^ sm3_rotate_left(w1[j - 3], 15))
              ^ sm3_rotate_left(w1[j - 13], 7) ^ w1[j - 6];
    }
    for (int j = 0; j < 64; ++j) {
        w2[j] = w1[j] ^ w1[j + 4];
    }
    // ABCDEFGH <- V(i)
    uint32_t a = sm3_data_get_uint32(v, 0);
    uint32_t b = sm3_data_get_uint32(v, 1);
    uint32_t c = sm3_data_get_uint32(v, 2);
    uint32_t d = sm3_data_get_uint32(v, 3);
    uint32_t e = sm3_data_get_uint32(v, 4);
    uint32_t f = sm3_data_get_uint32(v, 5);
    uint32_t g = sm3_data_get_uint32(v, 6);
    uint32_t h = sm3_data_get_uint32(v, 7);
    for (int j = 0; j < 64; ++j) {
        ss1 = sm3_rotate_left(
            sm3_rotate_left(a, 12)
            + e + sm3_rotate_left(j <= 15 ? SM3_T0 : SM3_T1, j),
            7
        );
        ss2 = ss1 ^ sm3_rotate_left(a, 12);
        tt1 = sm3_ff(a, b, c, j) + d + ss2 + w2[j];
        tt2 = sm3_gg(e, f, g, j) + h + ss1 + w1[j];
        d = c;
        c = sm3_rotate_left(b, 9);
        b = a;
        a = tt1;
        h = g;
        g = sm3_rotate_left(f, 19);
        f = e;
        e = sm3_p0(tt2);
    }
    // V(i+1) <- ABCDEFGH ^ V(i)
    sm3_data_xor_uint32(v, a, 0);
    sm3_data_xor_uint32(v, b, 1);
    sm3_data_xor_uint32(v, c, 2);
    sm3_data_xor_uint32(v, d, 3);
    sm3_data_xor_uint32(v, e, 4);
    sm3_data_xor_uint32(v, f, 5);
    sm3_data_xor_uint32(v, g, 6);
    sm3_data_xor_uint32(v, h, 7);
}

// SM3 Init context
void sm3_stream_start(sm3_ctx_t *ctx) {
    ctx->total_bits = 0;
    ctx->segment_bits = 0;
    memcpy(ctx->hash, SM3_IV, 32); // 256 bits
    memset(ctx->segment, 0, SM3_BIT_GROUP_SIZE / 8);   // 512 bits
}

// SM3 Update message segment
void sm3_stream_update(
    sm3_ctx_t *ctx,
    const uint8_t segment[],
    uint32_t segment_bits
) {
    assert(segment_bits <= SM3_BIT_GROUP_SIZE);
    // printf("SM3 Update %d bits\n", segment_bits);
    // sm3_show_bytes(segment, 64);
    if (ctx->segment_bits > 0) {
        // perform one step sm3_cf
        ctx->total_bits += ctx->segment_bits;
        sm3_cf(ctx->hash, ctx->segment);
    }
    // store last segment
    ctx->segment_bits = segment_bits;
    uint32_t segment_bytes = (segment_bits + 7) / 8;
    for (int i = 0; i < SM3_BIT_GROUP_SIZE / 8; ++i) {
        if (i < segment_bytes) {
            ctx->segment[i] = segment[i];
        } else {
            ctx->segment[i] = 0;
        }
    }
}

// SM3 Finalize hash result
void sm3_stream_finish(sm3_ctx_t *ctx, uint8_t hash[32]) {
    // process last segment data
    ctx->total_bits += ctx->segment_bits;
    if (ctx->total_bits > 0) {
        if (ctx->segment_bits % SM3_BIT_GROUP_SIZE == 0) {
            sm3_cf(ctx->hash, ctx->segment);
            memset(ctx->segment, 0, 64);
            ctx->segment[0] = 0x80; // append 0b10000000
        } else {
            // append 1
            ctx->segment[ctx->segment_bits / 8] |= 1 << (7 - ctx->segment_bits % 8);
            if (ctx->segment_bits + 1 + 64 > SM3_BIT_GROUP_SIZE) {
                sm3_cf(ctx->hash, ctx->segment);
                memset(ctx->segment, 0, 64);
            }
        }
        // last 8 bytes = message length
        uint64_t l = ctx->total_bits;
        for (int i = 0; i < 8; ++i, l /= 256) {
            ctx->segment[63 - i] = l & 0xFF;
        }
        // run the last segment
        sm3_cf(ctx->hash, ctx->segment);
    } else {
        memset(ctx->hash, 0, 32);
    }
    // Copy result to output
    for (int i = 0; i < 32; ++i) {
        hash[i] = ctx->hash[i];
    }
}

// SM3 SYNC CALL
void sm3_hash(uint8_t *hash, const uint8_t *message_ptr, uint64_t message_bits) {
    sm3_ctx_t ctx;
    sm3_stream_start(&ctx);
    // Group into 512 bits
    for (
            uint64_t i = 0, byte_offset = 0;
            i < message_bits;
            i += SM3_BIT_GROUP_SIZE, byte_offset += SM3_BIT_GROUP_SIZE / 8
            ) {
        if (message_bits - i >= SM3_BIT_GROUP_SIZE) {
            sm3_stream_update(&ctx, message_ptr + byte_offset, SM3_BIT_GROUP_SIZE);
        } else {
            sm3_stream_update(&ctx, message_ptr + byte_offset, message_bits - i);
        }
    }
    sm3_stream_finish(&ctx, hash);
}







